import os
import sys
import zlib
import struct


MAGIC_NUMBER = "\x55\xAA\x5A\xA5"
BOOT_MAGIC_NUMBER = "\x41\x4e\x44\x52\x4f\x49\x44\x21"
HEAD_LEN = 0x64
GZIP_MAGIC_NUMBER = "\x1F\x8B"
SYSTEM_IMG_MAGIC = '\x3a\xff\x26\xed'


def ungzip(gzip_path, ext_path):
    FTEXT, FHCRC, FEXTRA, FNAME, FCOMMENT = 1, 2, 4, 8, 16
    with open(gzip_path, "rb") as f:
        magic = f.read(2)
        if magic != '\037\213':
            return False

        method = ord(f.read(1))
        if method != 8:
            return False

        flag = ord(f.read(1))
        mtime = struct.unpack("<I", f.read(4))[0]
        f.read(2)

        if flag & FEXTRA:
            # Read & discard the extra field, if present
            xlen = ord(f.read(1))
            xlen = xlen + 256*ord(f.read(1))
            f.read(xlen)
        if flag & FNAME:
            # Read and discard a null-terminated string containing the filename
            while True:
                s = f.read(1)
                if not s or s=='\000':
                    break
        if flag & FCOMMENT:
            # Read and discard a null-terminated string containing a comment
            while True:
                s = f.read(1)
                if not s or s=='\000':
                    break
        if flag & FHCRC:
            f.read(2)     # Read & discard the 16-bit header CRC

        decompress = zlib.decompressobj(-zlib.MAX_WBITS)
        dec_data = decompress.decompress(f.read())
        if len(dec_data) == 0:
            return False

        with open(ext_path, "wb") as wf:
            wf.write(dec_data)
            return True

    return False

class Partition(object):
    def __init__(self, data, offset):
        if data[:4] != MAGIC_NUMBER:
            self.init = False
            return

        self.offset = offset
        self.magic_number = struct.unpack_from("I", data, 0)[0]
        self.head_size = struct.unpack_from("I", data, 4)[0]
        self.unknown4byte_1 = struct.unpack_from("I", data, 8)[0]
        self.unknown8byte = struct.unpack_from("8s", data, 0xc)[0]
        self.unknown4byte_2 = struct.unpack_from("I", data, 0x14)[0]
        self.data_size = struct.unpack_from("I", data, 0x18)[0]
        self.buid_date = struct.unpack_from("16s", data, 0x1c)[0].strip("\x00")
        self.buid_time = struct.unpack_from("16s", data, 0x2c)[0].strip("\x00")
        self.partition_name = struct.unpack_from("16s", data, 0x3c)[0].strip("\x00")
        self.null16byte = struct.unpack_from("16s", data, 0x4c)[0]
        self.head_crc16 = struct.unpack_from("h", data, 0x5c)[0]
        self.unknown4byte_3 = struct.unpack_from("I", data, 0x5e)[0]
        self.partition_crc16 = struct.unpack_from("h", data, 0x62)[0]

        self.init = True
        #self.show_info()

    def is_init(self):
        return self.init


    def show_info(self):
        print hex(self.magic_number)
        print hex(self.head_size)
        print hex(self.unknown4byte_1)
        print self.unknown8byte
        print hex(self.unknown4byte_2)
        print hex(self.data_size)
        print self.buid_date
        print self.buid_time
        print self.partition_name
        print self.null16byte
        print hex(self.head_crc16)
        print hex(self.unknown4byte_3)
        print hex(self.partition_crc16)



def is_update_app(path):
    offset = 0
    with open(path, "rb") as f:
        while True:
            if f.read(1) == MAGIC_NUMBER[0]:
                if f.read(3) == MAGIC_NUMBER[1:]:
                    f.seek(-4, os.SEEK_CUR)
                    offset = f.tell()
                    break;
                else:
                    f.seek(-3, os.SEEK_CUR)
        if offset == 0x5c:
            return True

    return False



def get_partitions(path):
    part_list = []
    with open(path, "rb") as f:
        flag = -1
        offset = 0
        while True:
            if f.read(1) == MAGIC_NUMBER[0]:
                if f.read(3) == MAGIC_NUMBER[1:]:
                    flag = 0
                    f.seek(-4, os.SEEK_CUR)
                    offset = f.tell()
                    break
                else:
                    f.seek(-3, os.SEEK_CUR)
        if flag != 0:
            return part_list

        while True:
            data = f.read(HEAD_LEN) 
            if len(data) != HEAD_LEN:
                break
            
            partition = Partition(data, offset)
            if not partition.is_init():
                break

            offset += partition.head_size + partition.data_size
            offset = ((offset + 4 - 1) & (~(4 - 1)))
            f.seek(offset, os.SEEK_SET)
            part_list.append(partition)
    return part_list

def extract_system(form_path, to_path, partitions):
    flag = -1
    system_offset = 0
    data_size = 0
    boot_begin = 0
    head_size = 0
    for i in partitions:
        if i.partition_name.lower() == "system":
            flag = 0
            system_offset = i.offset 
            data_size = i.data_size
            head_size = i.head_size
            break

    if flag != 0:
        return False

    print("system_offset = 0x%x, data_size = 0x%x" % (system_offset, data_size))

    with open(form_path, "rb") as f:
        f.seek(system_offset, os.SEEK_SET)
        system_data = f.read(data_size)
        system_img_offset = system_data.find(SYSTEM_IMG_MAGIC)
        if system_img_offset != -1:
            print("system_img_offset = 0x%x" % system_img_offset)
            system_offset = system_offset + system_img_offset

        system_size = data_size + head_size - system_img_offset

        f.seek(system_offset, os.SEEK_SET)
        system_data = f.read(system_size)
        if len(system_data) != system_size:
            return False

        with open(to_path, "wb") as wf:
            wf.write(system_data)
            return True

    return False

def extract_boot(form_path, to_path, partitions):
    flag = -1
    boot_offset = 0
    data_size = 0
    boot_begin = 0
    head_size = 0
    for i in partitions:
        if i.partition_name.lower() == "boot":
            flag = 0
            boot_offset = i.offset 
            data_size = i.data_size
            head_size = i.head_size
            break

    if flag != 0:
        return False

    flag = -1
    with open(form_path, "rb") as f:
        f.seek(boot_offset, os.SEEK_SET)
        while True:
            if f.read(1) == BOOT_MAGIC_NUMBER[0]:
                if f.read(7) == BOOT_MAGIC_NUMBER[1:]:
                    flag = 0
                    f.seek(-8, os.SEEK_CUR)
                    boot_begin = f.tell()
                    boot_size = (data_size + head_size) - (boot_begin - boot_offset)

                    print "boot_size = %d, data_size = %d, head_size = %d, boot_begin = %d, boot_offset = %d\n" %(boot_size, data_size, \
                    head_size, boot_begin, boot_offset)

                    break
                else:
                    f.seek(-7, os.SEEK_CUR)

        if flag != 0:
            return False

        boot_data = f.read(boot_size)
        if len(boot_data) != boot_size:
            return False

        with open(to_path, "wb") as wf:
            wf.write(boot_data)
            return True
    return False


def extract_zimage(boot_path, zimage_path):
    with open(boot_path, "rb") as f:
        if f.read(8) != BOOT_MAGIC_NUMBER:
            return False
        kernel_size = struct.unpack("I", f.read(4))[0]

        f.seek(0x24, os.SEEK_SET)
        page_size = struct.unpack("I", f.read(4))[0]

        f.seek(page_size, os.SEEK_SET)
        data = f.read(kernel_size)
        if len(data) != kernel_size:
            return False
        
        with open(zimage_path, "wb") as wf:
            wf.write(data)
            return True
    return False


def extract_kernel(zimage_path, code_path):
    with open(zimage_path, "rb") as f:
        data = f.read(2)
        if data[:2] == GZIP_MAGIC_NUMBER:
            return ungzip(zimage_path, code_path)
        else:
            f.seek(0, os.SEEK_SET)
            with open(code_path, "wb") as wf:
                wf.write(f.read())
            return True
    return False


def usage():
    info = "usage: extract_boot <update.app path> [extract path]"
    print(info)


def main():
    if not len(sys.argv) > 1 and len(sys.argv) <= 3:
        usage()
        return False

    update_path = sys.argv[1]
    if not os.path.exists(update_path):
        print("update.app do not exists")
        return False

    extract_path = os.path.join(os.path.dirname(update_path), "extract")
    if len(sys.argv) == 3:
        extract_path = sys.argv[2]

    if not os.path.exists(extract_path):
        os.mkdir(extract_path)

    boot_path = os.path.join(extract_path, "boot.img")
    zimage_path = os.path.join(extract_path, "zimage")
    code_path = os.path.join(extract_path, "code")
    system_path = os.path.join(extract_path, "system.img")
    if not is_update_app(update_path):
        print("do not update.app file")
        return False

    print("extract boot.img...")
    part_list = get_partitions(update_path)
    if not extract_boot(update_path, boot_path, part_list):
        print("extract boot.img fail")
        return False

    print("extract system.img ...")
    if not extract_system(update_path, system_path, part_list):
        print("extract system.img fail")
        return False

    print("extract zimage...")
    if not extract_zimage(boot_path, zimage_path):
        print("extract zimage fail")
        return False

    print("extract kernel...")
    if not extract_kernel(zimage_path, code_path):
        print("extract kernel fail")
        return False

    print("success")
    return True 


if __name__ == "__main__":
    #sys.argv.append(r"C:\Users\lihanghang\Desktop\xxx\UPDATE.APP")
    #sys.argv.append(r"C:\Users\lihanghang\Desktop\xxx")
    main()
